package com.atguigu.day05;

import com.atguigu.utils.ClickEvent;
import com.atguigu.utils.ClickSource;
import com.atguigu.utils.UrlViewCountPerWindow;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.List;

// 使用KeyedProcessFunction实现全窗口聚合函数
// 只需要将MapState的VALUE类型改成列表就可以了
public class Example4 {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment
                .getExecutionEnvironment();
        env.setParallelism(1);

        env
                .addSource(new ClickSource())
                .keyBy(r -> r.url)
                .process(new FakeWindow(5000L))
                .print();

        env.execute();
    }

    public static class FakeWindow extends
            KeyedProcessFunction<String, ClickEvent, UrlViewCountPerWindow> {
        private Long windowSize;

        public FakeWindow(Long windowSize) {
            this.windowSize = windowSize;
        }

        // KEY：窗口开始时间
        // VALUE：窗口开始时间KEY对应的窗口中的所有元素的LIST
        private MapState<Long, List<ClickEvent>> mapState;

        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);
            mapState = getRuntimeContext()
                    .getMapState(
                            new MapStateDescriptor<Long, List<ClickEvent>>(
                                    "map-state",
                                    Types.LONG,
                                    Types.LIST(Types.POJO(ClickEvent.class))
                            )
                    );
        }

        @Override
        public void processElement(ClickEvent value, Context ctx, Collector<UrlViewCountPerWindow> out) throws Exception {
            long currTs = ctx.timerService().currentProcessingTime();
            long windowStart = currTs - currTs % windowSize;
            long windowEnd = windowStart + windowSize;

            if (!mapState.contains(windowStart)) {
                ArrayList<ClickEvent> elements = new ArrayList<>();
                elements.add(value);
                mapState.put(windowStart, elements);
            } else {
                List<ClickEvent> elements = mapState.get(windowStart);
                elements.add(value);
                mapState.put(
                        windowStart,
                        elements
                );
            }

            ctx.timerService().registerProcessingTimeTimer(
                    windowEnd - 1L
            );
        }

        @Override
        public void onTimer(long timestamp, OnTimerContext ctx, Collector<UrlViewCountPerWindow> out) throws Exception {
            super.onTimer(timestamp, ctx, out);
            long windowEnd = timestamp + 1L;
            long windowStart = windowEnd - windowSize;
            long count = 0L;
            for (ClickEvent e : mapState.get(windowStart)) count++;
            mapState.remove(windowStart);
            String url = ctx.getCurrentKey();
            out.collect(new UrlViewCountPerWindow(
                    url,
                    count,
                    windowStart,
                    windowEnd
            ));
        }
    }

}
